# -*- coding: utf-8 -*-
import gym
from gym import spaces
import numpy as np
from net.msg_handle import msg_handle
from net.msg_class import *
import pandas as pd
from bdtime import Time
import logging
from bdtime import log_config_dc
logging.basicConfig(**log_config_dc.complex)
from queue import Queue
import threading
lock = threading.Lock()


class Car2dEnv(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'],
        'video.frames_per_second': 2
    }
    metadata['render_modes'] = metadata['render.modes']     # 避免报错

    def __init__(
            self,
            size=10,  # 地图大小
    ):
        self.target_x = 0
        self.target_y = 0
        self.size = size

        self.action_space = spaces.Discrete(5)  # 0, 1, 2，3，4: 不动，上下左右
        max_size = 1000.0
        # self.observation_space = spaces.Box(np.array([-self.size, -self.size]), np.array([self.size, self.size]))
        self.observation_space = spaces.Box(
            np.array([-max_size, -max_size, -max_size, -max_size], dtype=np.float32),
            np.array([max_size, max_size, max_size, max_size], dtype=np.float32)
        )

        self.state_msg_dc_ls = []
        self.state_msg_df = pd.DataFrame([])

        self.state_msg_queue = Queue()
        self.state_i_ls = []

        self.i_action = 0      # 需要一个i来对应起`发送的MsgAction`和从`client收到的MsgState`

        self.tt = Time()
        self.exceed_time = 0.3  # 太短容易造成tcp堵塞

    def set_exceed_time(self, exceed_time):
        self.exceed_time = exceed_time

    def get_formatted_ret(self, res):
        observation, reward, done, info = None, None, None, None
        if res is not None:
            try:
                _state_df = res.get('state')
                if (not isinstance(_state_df, str)) and len(_state_df) > 1:
                    _state_df = _state_df.iloc[0]
                state = json.loads(_state_df)
            except Exception as e:
                raise e
            _observation_dc, reward, done, info = state.get('obsevation'), state.get('reward'), state.get('done'), state.get('info'),
            observation: list = _observation_dc.get('agent_pos')
            observation.extend(_observation_dc.get('destination_pos'))
            observation = np.array(observation).reshape(len(observation), )

        ret = observation, reward, done, info
        return ret

    # Env: step(self, action: ActType) -> Tuple[ObsType, float, bool, dict]:
    def step(self, action):
        res = self._send_action_msg_and_wait_for_observation(action)
        ret = self.get_formatted_ret(res)
        return ret

    def on_msg_test(self, msg_dc):
        print('--- on_msg_test:', msg_dc)

        try:
            msg = MsgTest(self.tt.get_current_beijing_time_str())
            msg_handle.send(msg)
        except Exception as e:
            logging.exception(f"*** 多线程运行`on_msg_test`时出错, error: {e}")

    def on_msg_state(self, msg_dc):
        i = msg_dc.pop('i')
        state = msg_dc.get('state')
        with lock:
            if i in self.state_msg_df:
                return
            dc = {
                'i': i,
                'state': state
            }
            df_i = pd.DataFrame.from_records([dc], index='i')
            self.state_msg_df = self.state_msg_df.append(df_i)

    def _send_action_msg_and_wait_for_observation(self, action):
        """
        step发送MsgAction后, 等待客户端消息返回结果`state(s_, r, done, info)`, 以此提取轨迹 trajectory(s, a, r, done, s_)
        """
        msg = MsgAction(int(action), self.i_action)
        msg_handle.send(msg)

        tt = Time()

        state = None

        tt.__init__()
        while tt.during(self.exceed_time):
            if not msg_handle.client_connection:
                break

            if self.i_action in self.state_msg_df.index:
                try:
                    state = self.state_msg_df.loc[self.i_action]
                except Exception as e:
                    print("\n\n****** state_msg_df.loc error!")
                    print('!' * 20)
                    print(e)
                    state = None
                    break
                # --- 超过200条后, 批量删除前100条
                total_state_msg = len(self.state_msg_df)
                drop_len = 100
                if total_state_msg > drop_len and total_state_msg % drop_len == 0:
                    self.state_msg_df = self.state_msg_df.drop(self.state_msg_df.index[:drop_len])

                break
            tt.sleep(0.0001)

        if not msg_handle.client_connection:
            print('*** env.step: 客户端链接已断开, stop learning.')
            return None

        if state is None:
            # msg_handle.close()
            print('\n\n****** Error: state 为 None!')
            print(f'tt.exceed({self.exceed_time}):', tt.exceed(self.exceed_time))
            print(f'state为空? msg: {msg}')
            print('*' * 50)

        self.i_action += 1
        return state

    def reset(self):
        res = self._send_action_msg_and_wait_for_observation(0)
        ret = self.get_formatted_ret(res)
        return ret


